#!/usr/bin/env python


# roc.py: given a matching score matrix (smaller is better) and a "flag matrix" that identifies
# correct match scores and incorrect match scores in the matching score matrix, compute the data on the ROC curve.
# The output is a CSV file.  The inputs are CSV files also.
#
# 2/13/08 PJF
#
# flag matrix matches the match matrix in layout.
# entries are as follows: 0 means ignore (use for header row and column),
# 1 means match, 2 means nonmatch
#

import sys
import getopt


if len(sys.argv) < 4:
	print 'Usage: roc [-r] matchmatrix.csv flagmatrix.csv outfile.csv'
	print '-r: assume higher scores are better matches (default: lower is better)'
	sys.exit(-1)

lowerbetter = True

oa,args = getopt.getopt(sys.argv[1:],"r")

for o,a in oa:
	if o == '-r': lowerbetter = False




# read matching matrix entries
print 'score matrix file: ',args[0]
fmatch = file(args[0],'r')

if not fmatch:
	print 'match file open failed.'
	sys.exit(-1)

# gobble up all of the lines in the file.
matchlines = fmatch.readlines()

fmatch.close()  

#parse matching matrix entries
print len(matchlines),'lines in matching score matrix'
scores = []

for l in matchlines:
	l = l.strip().split(',')
	scores.append(l)

print len(scores),'rows in matching score matrix'

# read flag matrix entries
fmask = file(args[1],'r')

if not fmask:
	print 'mask file open failed.'
	sys.exit(-1)

# scarf up the lines in the mask matrix.
masklines = fmask.readlines()
fmask.close()  

#parse mask matrix entries
print len(masklines),'lines in mask matrix'
flags = []

for l in masklines:
	l = l.strip().split(',')
	print 'maskline: ',l
	flags.append(l)

print len(flags),'rows in flag matrix'

# sanity checks

if len(scores) != len(flags): 
	print 'ERROR. length mismatch between scores and flags.'
	sys.exit(-1)

#
# use flags to assemble list of match and nonmatch scores
#

matchscores = []
nonmatchscores = []

for i in xrange(len(flags)):
	f = flags[i]
	s = scores[i]
	if len(f) != len(s):
		print 'ERROR. length mismatch between flag and score row ',i
		print 'flag: ',len(f),' score:',len(s)
		sys.exit(-1)
	while(f):
		ftmp = f.pop()
		stmp = s.pop()
		if ftmp == '1':
			matchscores.append(float(stmp))
		elif ftmp == '2':
			nonmatchscores.append(float(stmp))
		else:
			print 'skipped ',stmp

print len(matchscores),' match scores extracted.'
print len(nonmatchscores),' nonmatch scores extracted.'

matchscores.sort()
nonmatchscores.sort()

#
# OK, here's the hard part. The decision of a match versus nonmatch is based on a threshold on
# matching score.  When a score is computed, we decide whether it is a match or not.
# Since the truth isn't known, we could have four cases: correct match, false match, correct nonmatch, false nonmatch.
#

scores = matchscores + nonmatchscores
scores.sort()

froc = file(args[2],'w')
froc.write('threshold,far,frr\n')

if lowerbetter:
	while scores:
		thr = scores.pop()
		nfa = 0
		nfr = 0
		for x in matchscores: # remember, lower is better
			if x > thr:
				nfr = nfr+1
		for x in nonmatchscores:
			if x <= thr:
				nfa = nfa+1
		far = 1.0*nfa/len(nonmatchscores)
		frr = 1.0*nfr/len(matchscores)
		froc.write('%g,%g,%g\n' % (thr,far,frr))
else:
	while scores:
		thr = scores.pop()
		nfa = 0
		nfr = 0
		for x in matchscores: # remember, higher is better
			if x < thr:
				nfr = nfr+1
		for x in nonmatchscores:
			if x >= thr:
				nfa = nfa+1
		far = 1.0*nfa/len(nonmatchscores)
		frr = 1.0*nfr/len(matchscores)
		froc.write('%g,%g,%g\n' % (thr,far,frr))
			
froc.close()
